from __future__ import (absolute_import, division, print_function, unicode_literals)
import datetime
import backtrader as bt
import akshare as ak
import pandas as pd


# https://blog.csdn.net/m0_46603114/article/details/105014498

def get_data(code):
    df = ak.stock_zh_a_daily(symbol=code, adjust='qfq')
    df.index = pd.to_datetime(df.date)
    df.sort_index(inplace=True)

    return df


class St(bt.Strategy):
    params = dict(
        buy_limit_percent=0.01,
        buy_valid_date=5,
        stoptype=bt.Order.StopTrail,
        trailamount=0.0,
        trailpercent=0.05,
        p_high_period=5,
        p_fast=5,
        p_slow=60,
    )

    def __init__(self):
        self.slowSMA = bt.ind.SMA(period=self.p.p_slow)
        self.fastSMA = bt.ind.SMA(period=self.p.p_fast)
        self.buy_con = bt.And(bt.ind.CrossUp(self.fastSMA, self.slowSMA),
                              # slowSMA == bt.ind.Highest(slowSMA, period = self.p.p_high_period, plot = False)
                              )
        self.order = None

    def next(self):
        print(
            f"今日：datetime {self.data.lines.datetime.date(0)} sma_slow {self.slowSMA[0]} sma_fast {self.fastSMA[0]}")
        print(
            f"今日：datetime {self.data.lines.datetime.date(0)} open {self.data.lines.open[0]} close {self.data.lines.close[0]} high {self.data.lines.high[0]} low {self.data.lines.low[0]}")
        # 无场内资产
        if not self.position:
            # 未提交买单
            if self.order is None:
                # 金叉到达了买点
                if self.buy_con:
                    # 计算订单有效期时间，如果超过有效期，股价仍未回踩，则放弃下买入订单
                    valid = self.data.datetime.date(0)
                    if self.p.buy_valid_date:
                        valid = valid + datetime.timedelta(days=self.p.buy_valid_date)
                    # 计算回踩后的买入价格
                    price = self.datas[0].close[0] * (1.0 - self.p.buy_limit_percent)
                    print(f"Buy order created: {self.datetime.date()}: close: {self.datas[0].close[0]} "
                          f"limit price: {price} valid: {valid}")
                    # 用有效时间及回踩买点提交买入订单
                    # 如果开盘价低于price，那么订单就会用开盘价被立即执行; 如果开盘价高于price，但是当日最低价低于price，
                    # 那么订单将会以price的价格被执行从Limit订单的描述可以看出，正好符合我们所说的股价回踩买入规则
                    self.order = self.buy(exectype=bt.Order.Limit, price=price, valid=valid)
                    # o = self.buy()
                    print('*' * 50)

        elif self.order is None:
            # 提交stoptrail订单
            self.order = self.sell(exectype=self.p.stoptype,
                                   trailamount=self.p.trailamount,
                                   trailpercent=self.p.trailpercent)
            if self.p.trailamount:
                tcheck = self.data.close - self.p.trailamount
            else:
                tcheck = self.data.close * (1.0 - self.p.trailpercent)
            print(f"Sell Block date: {self.datetime.date()}: "
                  f"close： {self.data.close[0]} "
                  f"Limit price: {self.order.created.price} check price {tcheck}")
            print('-' * 10)
        else:
            if self.p.trailamount:
                tcheck = self.data.close - self.p.trailamount
            else:
                tcheck = self.data.close * (1.0 - self.p.trailpercent)
            print(f"update limit price: {self.datetime.date()} close： {self.data.close[0]} "
                  f"Limit price: {self.order.created.price} / check price {tcheck}")

    def notify_order(self, order):
        # 未被处理的订单
        if order.status in [order.Submitted, order.Accepted]:
            return
        # 已经处理的订单
        if order.status in [order.Completed, order.Canceled, order.Margin]:
            if order.isbuy():
                self.log(f'code={order.data._name} BUY CREATE TIME: {bt.num2date(order.created.dt)},'
                         f' EXECUTED TIME: {bt.num2date(order.executed.dt)} EXECUTED PRICE: {order.executed.price}')

            else:  # Sell
                self.log(f'code={order.data._name} SELL CREATE TIME: {bt.num2date(order.created.dt)},'
                         f' EXECUTED TIME: {bt.num2date(order.executed.dt)} EXECUTED PRICE: {order.executed.price}')

            self.bar_executed = len(self)

        # Write down: no pending order
        self.order = None

    def log(self, txt, dt=None, doprint=True):
        if doprint:
            dt = dt or self.datas[0].datetime.date(0)
            # print(f'{dt.isoformat()},{txt}')
            print(txt)


codes = ['sz000001', 'sz000002', 'sh603999']
code = codes[2]

st_date = datetime.datetime(2022, 1, 1)
ed_date = datetime.datetime(2023, 6, 16)
data = get_data(code)

cerebro = bt.Cerebro()  # 创建cerebro
datafeed = bt.feeds.PandasData(dataname=data, fromdate=st_date, todate=ed_date)
cerebro.adddata(datafeed, name=code)
# 设置启动资金
cerebro.broker.setcash(100000.0)
# 设置交易单位大小
cerebro.addsizer(bt.sizers.FixedSize, stake=1000)
# 设置佣金为千分之一
cerebro.broker.setcommission(commission=0.001)
cerebro.addstrategy(St)  # 添加策略
cerebro.run()  # 遍历所有数据
# 打印最后结果
print('Final Portfolio Value: %.2f' % cerebro.broker.getvalue())
cerebro.plot(style='candlestick')  # 绘图
